# 控制模式2 passive模式

import mujoco
import math
import control as ct
import numpy as np
import mujoco.viewer as viewer
from os.path import dirname, join
import time



# 加载模型
model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene.xml')) 
data = mujoco.MjData(model)

key_id = model.key("home").id
mujoco.mj_resetDataKeyframe(model, data, key_id)

damping: float = 1e-4
diag = damping * np.eye(6)

site_id = model.site('site_Hand').id
base_id = model.body('BigArm1').id
goal_pos = [0.5, 0.5, 0.5]
end_effector_id = model.body('Hand').id
jac = np.zeros((6, model.nv))

mocap_id = model.body("target").mocapid[0]

# 这里令位姿是一个6x1的向量，1-3是rot，4-6是pos
error = np.zeros(6)
error_ori = error[3:]
error_pos = error[:3]



# 定义site姿态四元数
site_quat = np.zeros(4)
site_quat_conj = np.zeros(4)
error_quat = np.zeros(4)




def impedanceControl(model,data):
    # print(f"sensordata is {data.sensordata}")
    # 将末端执行器当前姿态的旋转矩阵转换为四元数形式（site_quat）
    mujoco.mju_mat2Quat(site_quat, data.site(site_id).xmat)

    # 计算四元数的共轭（表示逆向旋转），用于后续误差计算
    mujoco.mju_negQuat(site_quat_conj, site_quat)

    # 计算目标姿态四元数与当前姿态四元数的旋转差（误差四元数）
    # 通过将目标四元数（mocap_quat）与当前四元数的共轭相乘得到相对旋转
    mujoco.mju_mulQuat(error_quat, data.mocap_quat[mocap_id], site_quat_conj)

    # 将四元数形式的旋转误差转换为3D角速度向量（error_ori）
    # 参数1.0表示时间步长，用于规范化角速度输出
    mujoco.mju_quat2Vel(error_ori, error_quat, 1.0)

    # 求Frame Jacobian，即从关节角度到执行器位姿的雅克比矩阵。
    # 具体格式为mujoco.mj_jacSite(model, data, jacp_expected, jacr_expected, site_id)
    mujoco.mj_jacSite(model, data, jac[:3], jac[3:], site_id)

    # 对于本程序中的error定义和J定义，我们令J @ v = error则可以实现IK，此外我们引入damping防止奇异
    
    
    v_star = np.linalg.solve(jac.T @ jac + diag, jac.T @ (error))

    kv = 1


    if np.linalg.norm(error) > 1e-2:
        adaptive_kv = kv * np.tanh(np.linalg.norm(error)/0.1)  # 自适应增益
        ctrl = adaptive_kv*v_star
    else:
        ctrl = np.zeros(model.nu)

    print(f"ctrl is {ctrl}")

    return ctrl



with viewer.launch_passive(model,data) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(model, viewer.cam)
    while viewer.is_running():
        data.ctrl = impedanceControl(model,data)  # apply control
        mujoco.mj_step(model, data)            # step xdot= f(x,u)
        mujoco.mj_forward(model, data)


        viewer.sync()
        time.sleep(model.opt.timestep)





        # # Mujoco目前不能直接输出SE3矩阵
        # site_pos = data.site_xpos[site_id]
        # site_rot = data.site_xmat[site_id].reshape(3,3)
        # SE3 = np.eye(4)
        # SE3[:3, :3] = site_rot  # 填充旋转矩阵部分
        # SE3[:3, 3] = site_pos   # 填充位置向量部分
        # print(f"site_pos = {site_pos}\nsite_rot = {site_rot}\n")
        # print(f"SE3 = {SE3}\n")